package com.zytech.zspider.task.impl;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.zytech.modules.lang.AggregateFunction;
import com.zytech.modules.lang.Langs;
import com.zytech.modules.lang.Requires;
import com.zytech.zspider.core.*;
import com.zytech.zspider.logger.Loggers;
import com.zytech.zspider.task.*;

import java.math.BigDecimal;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.RecursiveTask;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author siuming
 */
public abstract class ForkJoinTaskExecutor implements TaskExecutor {

    // default blocking factor.
    protected static final float DEFAULT_BLOCKING_FACTOR = 0.5f;

    private final ForkJoinPool pool;
    private final AtomicReference<TaskExecutor.Status> statusRef;

    private final TaskGenerator taskGenerator;
    private final TaskExpressionEvaluator expressionEvaluator;

    /**
     * @param evaluator
     */
    protected ForkJoinTaskExecutor(String[] externalScripts) {
        this(DEFAULT_BLOCKING_FACTOR, externalScripts);
    }

    /**
     * @param blockingFactor
     * @param evaluator
     */
    protected ForkJoinTaskExecutor(float blockingFactor, String[] externalScripts) {
        Requires.isTrue(blockingFactor > 0 && blockingFactor < 1, "blockingFactor must between (0,1).");

        this.expressionEvaluator = new TaskExpressionEvaluatorImpl(externalScripts);
        this.taskGenerator = new TaskGeneratorImpl(expressionEvaluator);

        BigDecimal cpus = new BigDecimal(Langs.mycpus());
        BigDecimal threads = cpus.divide(new BigDecimal(1).subtract(new BigDecimal(blockingFactor)), BigDecimal.ROUND_HALF_UP);
        this.pool = new ForkJoinPool(threads.intValue());
        this.statusRef = new AtomicReference<Status>(Status.OPEN);

    }

    @Override
    public final TaskResult submitTask(TaskDefinition task) {
        Requires.notNull(task, "task must not be null.");
        ensureSelfIsOpen();
        try {

            return getPool()
                    .submit(new TaskAction(TaskDefinition.EMPTY, TaskContext.EMPTY, TaskResult.EMPTY, task))
                    .get();
        } catch (Exception e) {
            throw handleTaskSubmitException(e);
        }
    }

    private void ensureSelfIsOpen() {
        if (getStatus() != Status.OPEN) {
            throw new IllegalStateException("TaskExecutor must not be shutdown.");
        }
    }

    @Override
    public final TaskResult retryTask(TaskDefinition task, TaskContext context) {
        Requires.notNull(task, "task must not be null.");
        Requires.notNull(context, "context must not be null.");
        return onTaskSubmit(task, context);
    }

    /**
     * @param e
     * @return
     */
    protected TaskException handleTaskSubmitException(Exception e) {
        if (e instanceof InterruptedException) {
            Thread.currentThread().interrupt();
        }

        Loggers.me().error(getClass(), "task submit occurs error.", e);
        return TaskExceptions.transform(e);
    }

    final class TaskAction extends RecursiveTask<TaskResult> {

        private static final long serialVersionUID = -8515582352072753455L;

        final TaskDefinition father;
        final TaskContext fatherContext;
        final TaskResult fatherResult;
        final TaskDefinition task;

        /**
         * @param father
         * @param fatherContext
         * @param fatherResult
         * @param task
         */
        TaskAction(TaskDefinition father, TaskContext fatherContext,
                   TaskResult fatherResult, TaskDefinition task) {
            this.father = father;
            this.fatherContext = fatherContext;
            this.fatherResult = fatherResult;
            this.task = task;
        }

        @Override
        protected TaskResult compute() {

            TaskContext context = newTaskContext();
            onTaskStart(task, context);
            TaskResult result = onTaskSubmit(task, context);
            TaskResult usingResult = onTaskResult(task, context, result);
            onTaskFinish(task, context, usingResult);

            List<TaskAction> actions = computeSubtasks(task, context, usingResult);
            return Langs.aggregate(usingResult, actions, new AggregateFunction<TaskResult, TaskAction>() {
                @Override
                public TaskResult apply(TaskResult result, TaskAction action) {
                    return result.combine(action.join());
                }
            });
        }

        private TaskContext newTaskContext() {
            TaskContext context = null != fatherContext ? fatherContext.copy() : TaskContext.of(task);
            context.putBinding(TaskContext.TASK, task);
            context.putBinding(TaskContext.FATHER_TASK, null == father ? father : father.copy());
            context.putBinding(TaskContext.FATHER_TASK_RESULT, fatherResult);
            return context;
        }
    }

    private List<TaskAction> computeSubtasks(TaskDefinition task,
                                             TaskContext context,
                                             TaskResult result) {

        if (task.isEmptySubtaskGenDefinitions() && task.isEmptySubtasks()) {
            return Collections.emptyList();
        }

        List<TaskDefinition> gens = generateSubtasks(task, context, result);
        List<TaskDefinition> subtasks = ImmutableList.
                <TaskDefinition>builder()
                .addAll(gens)
                .addAll(task.getSubtasks())
                .build();

        Loggers.me().info(getClass(), String.format("submit %d subtasks for %s", subtasks.size(), task));
        List<TaskAction> actions = Lists.newArrayList();
        for (TaskDefinition subtask : subtasks) {
            subtask.diffSettings(task.getSettings()).diffHeaders(task.getHeaders());
            TaskAction action = (TaskAction) new TaskAction(task, context, result, subtask).fork();
            actions.add(action);
        }
        return actions;
    }

    /**
     * @param task
     * @param context
     * @param result
     * @return
     */
    protected final List<TaskDefinition> generateSubtasks(TaskDefinition task,
                                                          TaskContext context,
                                                          TaskResult result) {
        if (task.isEmptySubtaskGenDefinitions()) {
            return Collections.emptyList();
        }

        List<TaskDefinition> results = Lists.newArrayList();
        for (TaskGenDefinition genDef : task.getSubtaskGenDefinitions()) {
            List<TaskDefinition> subtasks = generateSubtaskByDef(genDef, context, result);
            results.addAll(subtasks);
        }
        return results;
    }

    private List<TaskDefinition> generateSubtaskByDef(TaskGenDefinition genDef, TaskContext context, TaskResult result) {
        try {

            context.putBinding(TaskContext.TASK_RESULT, result);
            return getTaskGenerator().generateTasks(context, genDef);
        } finally {
            context.putBinding(TaskContext.TASK_RESULT, TaskResult.EMPTY);
        }
    }

    /**
     * @param task
     * @param context
     */
    protected void onTaskStart(TaskDefinition task, TaskContext context) {
        if (task.isEmptyStartExpr()) {
            return;
        }

        Loggers.me().info(getClass(), "call start expr '''{}''' for {}.", task.getStartExpr(), task);
        getExpressionEvaluator().eval(task.getStartExpr(), context);
    }

    /**
     * @param task
     * @param context
     * @param result
     * @return
     */
    protected TaskResult onTaskResult(TaskDefinition task, TaskContext context, TaskResult result) {
        // TODO
        TaskResultHandler handler = new TaskResultHandlerPipeline(
                new TaskResultRetryHandler(),
                new TaskResultResetHandler()
        );
        return handler.onResult(this, task, context, result);
    }

    /**
     * @param task
     * @param context
     * @param result
     */
    protected void onTaskFinish(TaskDefinition task, TaskContext context, TaskResult result) {
        if (task.isEmptyFinishExpr()) {
            return;
        }

        try {

            Loggers.me().info(getClass(), "call finish expr '''{}''' for {}.", task.getFinishExpr(), task);
            context.putBinding(TaskContext.TASK_RESULT, result);
            getExpressionEvaluator().eval(task.getFinishExpr(), context);
        } finally {
            context.putBinding(TaskContext.TASK_RESULT, TaskResult.EMPTY);
        }
    }

    /**
     * @param current
     * @param context
     * @return
     */
    protected abstract TaskResult onTaskSubmit(TaskDefinition current, TaskContext context);

    public final void shutdown() {
        try {

            if (statusRef.compareAndSet(Status.OPEN, Status.SHUTDOWNING)) {
                Loggers.me().info(getClass(), "shutting down...");
                doShutdown();
            }
        } finally {
            if (statusRef.compareAndSet(Status.SHUTDOWNING, Status.SHUTDOWN)) {
                Loggers.me().info(getClass(), "shut down at ts[{}].", System.currentTimeMillis());
            }

        }
    }

    protected void doShutdown() {
        if (!getPool().isShutdown()) {
            getPool().shutdown();
        }
    }

    /**
     * @return
     */
    protected final Status getStatus() {
        return statusRef.get();
    }

    /**
     * @return
     */
    protected final ForkJoinPool getPool() {
        return pool;
    }

    /**
     * @return
     */
    protected final TaskGenerator getTaskGenerator() {
        return taskGenerator;
    }

    @Override
    public final TaskExpressionEvaluator getExpressionEvaluator() {
        return expressionEvaluator;
    }
}
